# 百度OCR API发票识别模块

import base64
import requests
import json
import os
import datetime


class InvoiceRecognizer:
    """
    使用百度OCR API进行增值税发票识别的类
    """

    def __init__(self, api_key='', secret_key=''):
        self.api_url = "https://aip.baidubce.com/rest/2.0/ocr/v1/vat_invoice"
        self.token_url = "https://aip.baidubce.com/oauth/2.0/token"
        self.api_key = api_key
        self.secret_key = secret_key
        self.access_token = None

        # 如果提供了API密钥，则获取token
        if api_key and secret_key:
            self.get_access_token()
        else:
            # 使用默认token（应在生产环境中替换为真实获取的token）
            self.access_token = '24.fb0e6430b275dca6033c8b6e9b5b0166.2592000.1744450324.282335-118025493'

    def get_access_token(self):
        """
        获取百度API的access_token
        """
        params = {
            'grant_type': 'client_credentials',
            'client_id': self.api_key,
            'client_secret': self.secret_key
        }

        response = requests.post(self.token_url, params=params)
        if response.status_code == 200:
            result = response.json()
            self.access_token = result.get('access_token')
            return self.access_token
        else:
            raise Exception(f"获取token失败: {response.text}")

    def recognize_invoice(self, image_path):
        """
        识别发票图片

        参数:
        image_path: 发票图片路径

        返回:
        解析后的发票信息字典
        """
        if not self.access_token:
            raise Exception("未设置access_token")

        if not os.path.exists(image_path):
            raise FileNotFoundError(f"找不到图片文件: {image_path}")

        # 读取图片文件并进行base64编码
        with open(image_path, 'rb') as f:
            image_data = base64.b64encode(f.read())

        # 构建请求参数
        params = {"image": image_data}
        request_url = f"{self.api_url}?access_token={self.access_token}"
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}

        # 发送请求
        response = requests.post(request_url, data=params, headers=headers)

        print(response.text)

        if response.status_code != 200:
            raise Exception(f"API请求失败: {response.status_code}, {response.text}")

        # 解析响应数据
        result = response.json()

        if 'error_code' in result:
            raise Exception(f"识别错误: {result.get('error_msg', '未知错误')}")

        # 解析发票要素
        return self.parse_invoice_data(result)

    def parse_invoice_data(self, response_json):
        """
        解析API返回的发票数据

        参数:
        response_json: API返回的JSON数据

        返回:
        解析后的发票信息字典
        """
        if 'words_result' not in response_json:
            raise Exception("API返回数据格式错误")

        result = response_json['words_result']

        # 提取所需要素
        invoice_elements = {
            "发票号码": result.get("InvoiceNum", ""),
            "开票日期": result.get("InvoiceDate", ""),
            "销售方名称": result.get("SellerName", ""),
            "购买方名称": result.get("PurchaserName", ""),
            "项目名称": result["CommodityName"][0]["word"] if result.get("CommodityName") else "",
            "不含税金额": result.get("TotalAmount", ""),
            "税额": result.get("TotalTax", ""),
            "含税金额": result.get("AmountInFiguers", "")
        }

        # 处理日期格式（从"2025年02月04日"转换为"2025-02-04"）
        if invoice_elements["开票日期"]:
            date_str = invoice_elements["开票日期"]
            date_str = date_str.replace('年', '-').replace('月', '-').replace('日', '')
            try:
                date_obj = datetime.datetime.strptime(date_str, "%Y-%m-%d")
                invoice_elements["开票日期"] = date_obj.strftime("%Y-%m-%d")
            except ValueError:
                # 如果日期格式无法解析，保留原始格式
                pass

        return invoice_elements


# 使用示例
if __name__ == "__main__":
    recognizer = InvoiceRecognizer()
    try:
        # 确保替换为你的发票图片路径
        invoice_data = recognizer.recognize_invoice("invoice1.jpg")
        print(json.dumps(invoice_data, ensure_ascii=False, indent=2))
    except Exception as e:
        print(f"识别失败: {str(e)}")